package 左哥算法.ch06二叉树;

import java.util.Comparator;

public class AVLTree<T> extends SearchTree<T> {
    int height = 1;

    public AVLTree(T val, Comparator<? super T> comparator) {
        super(val, comparator);
    }

    public AVLTree(T val) {
        super(val);
    }

    /**
     * @param val：被添加的值
     * @return：返回新的头节点
     */
    @Override
    public SearchTree<T> add(T val) {
        AVLTree<T> node = new AVLTree<>(val);
        super.add(node);
        return avlCheck(node.parent);   //返回新的头节点
    }

    @Override
    public SearchTree<T> remove(T key) {
        SearchTree<T> node = find(key);
        AVLTree<T> newRoot = (AVLTree<T>) super.remove(node);
        return newRoot.avlCheck(node.parent);
    }

    /**
     * 从index节点检查到curr节点
     *
     * @param checkNode：被检查的节点
     * @return：如果发生旋转，则返回新的头节点，否则返回原来的头节点
     */
    public SearchTree<T> avlCheck(SearchTree<T> checkNode) {
        updateHeight(checkNode);
        SearchTree<T> curr = checkNode;
        SearchTree<T> pre = curr;
        while (curr != null) {
            SearchTree<T> parent = curr.parent;
            int leftHeight = curr.left == null ? 0 : ((AVLTree<T>) curr.left).height;
            int rightHeight = curr.right == null ? 0 : ((AVLTree<T>) curr.right).height;
            if (leftHeight - rightHeight > 1) {  //左树导致的不平衡
                int L = curr.left.left == null ? 0 : ((AVLTree<T>) curr.left.left).height;
                int R = curr.left.right == null ? 0 : ((AVLTree<T>) curr.left.right).height;
                if (L > R) {   //LL型
                    curr = breakByLL(curr);
                    updateHeight(curr.right);
                } else {     //LR型
                    curr = breakByLR(curr);
                    updateHeight(curr.left);
                    updateHeight(curr.right);
                }
            }
            if (rightHeight - leftHeight > 1) {    //右树导致的不平衡
                int L = curr.right.left == null ? 0 : ((AVLTree<T>) curr.right.left).height;
                int R = curr.right.right == null ? 0 : ((AVLTree<T>) curr.right.right).height;
                if (R > L) {     //RR型
                    curr = breakByRR(curr);
                    updateHeight(curr.left);
                } else {     //RL型
                    curr = breakByRL(curr);
                    updateHeight(curr.right);
                    updateHeight(curr.left);
                }
            }
            pre = curr;
            curr = curr.parent;
        }
        return pre;
    }

    /**
     * LL类型的不平衡
     *
     * @param node：不平衡的头节点
     * @return：返回新的头节点
     */
    public SearchTree<T> breakByLL(SearchTree<T> node) {
        return node.rightSpin();
    }
    public SearchTree<T> breakByLR(SearchTree<T> node) {
        node.left = node.left.leftSpin();
        return breakByLL(node);
    }
    public SearchTree<T> breakByRR(SearchTree<T> node) {
        return node.leftSpin();
    }
    public SearchTree<T> breakByRL(SearchTree<T> node) {
        node.right = node.right.rightSpin();
        return breakByRR(node);
    }

    /**
     * @param node：从该节点去维护高度
     */
    public void updateHeight(SearchTree<T> node) {
        AVLTree<T> curr = (AVLTree<T>) node;
        while (curr != null) {
            //childHeight：孩子节点中较高的一个
            int childHeight = Math.max(curr.left == null ? 0 : ((AVLTree<T>) curr.left).height, curr.right == null ? 0 : ((AVLTree<T>) curr.right).height);
            curr.height = childHeight + 1;  //更新当前高度
            curr = (AVLTree<T>) curr.parent;
        }
    }
}